package algorithm_demo.demo04.graph;

import java.util.Comparator;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

/**
 * 最小生成树：Prim算法
 *
 * @author Api
 * @date 2023/2/19 21:53
 */
public class Code05_Prim {

    public static Set<Edge> primMST(Graph graph) {
        //解锁的边进入小根堆
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(Comparator.comparingInt(o -> o.weight));
        //哪些点被解锁出来了
        Set<Node> nodeSet = new HashSet<>();
        //已经考虑过的边不要重复考虑
        Set<Edge> edgeSet = new HashSet<>();

        Set<Edge> result = new HashSet<>();//一次挑选的边在result里
        for (Node node : graph.nodes.values()) {//随便挑了一个点
            //node是开始点
            if (!nodeSet.contains(node)) {
                nodeSet.add(node);
                for (Edge edge : node.edges) {//由一个点，解锁所有相连的边
                    if (!edgeSet.contains(edge)) {
                        edgeSet.add(edge);
                        priorityQueue.add(edge);
                    }
                }
                while (!priorityQueue.isEmpty()) {
                    Edge edge = priorityQueue.poll();//弹出解锁的边中，最小的边
                    Node toNode = edge.to;//可能的一个新的点
                    if (!nodeSet.contains(toNode)) {
                        nodeSet.add(toNode);
                        result.add(edge);
                        for (Edge nextEdge : toNode.edges) {
                            if (!edgeSet.contains(nextEdge)){
                                edgeSet.add(nextEdge);
                                priorityQueue.add(nextEdge);
                            }
                        }
                    }
                }
            }
        }
        return result;
    }
}
